Skip to content

Rename notebooks/ → livebooks/ - #178

Merged
ausimian merged 5 commits into
feat/expr-compilerfrom
feat/expr-compiler-livebooks-rename
Jun 6, 2026
Merged

Rename notebooks/ → livebooks/#178
ausimian merged 5 commits into
feat/expr-compilerfrom
feat/expr-compiler-livebooks-rename

Conversation

@ausimian

@ausimian ausimian commented Jun 6, 2026

Copy link
Copy Markdown
Owner

The example documents are Livebooks (.livemd), so name the directory to
match. Renames notebooks/livebooks/ (all 8 files) and updates every
current-state reference: mix.exs extras:/groups_for_extras,
scripts/test-livebooks.sh, README.md, ROADMAP.md. CHANGELOG.md /
PLAN.md keep their historical notebooks/ mentions.

Docs/rename only — independent of the rest of the stack (no shared files);
sequenced here only to keep the branch chain linear.


PR 3/5. Stacked on …-take-along-axis; lands in feat/expr-compiler.

ausimian added 5 commits June 6, 2026 13:13
Add a `mode_test/2` macro to `Emily.ConformanceHelper` that emits each
conformance test in three lanes from a single body, varying only the
bound `predict_opts`:

  * eval     — `[]`, the evaluator path (unchanged behaviour)
  * native   — `native: true, native_fallback: :raise`
  * fusion   — additionally `native_compiled: true` (mx::compile)

The native and fusion lanes are tagged `:native` and `:native_compiled`
on top of the module's `:conformance`, so `--only conformance` runs all
three while `--only native` / `--only native_compiled` select one lane.
Both compiler lanes resolve the same HuggingFace repos as the evaluator
lane, so the checkpoint download is amortised across them.

Convert the six tiny-random suites (DistilBERT, ViT, Whisper, ModernBert,
NomicBert, SmolLm3). `Axon.predict`-based tests thread `predict_opts`
into the call; `Axon.build`-based smoke tests build `init_fn` on the
evaluator and only `predict_fn` under `predict_opts`, so the native gate
covers the forward pass rather than random param init. The serving and
fast-kernels tests stay eval-only.

`native_fallback: :raise` makes the native lanes a no-fallback op-coverage
gate: every architecture lowers fully, and the fusion lane holds at the
existing 1e-4 reference tolerance (mx::compile only reassociates f32, so
its drift is far below the approximate mx::fast::* kernels).
Extend `mode_test/2` to `mode_test/3` with a `lane_tags:` option and run
the `*_full` forward-pass conformance suites (ViT-Base, Whisper-tiny)
through the native and fusion compilers alongside the evaluator.

The full suites pass `lane_tags: false`, so their native and fusion
lanes are emitted without the cross-cutting `:native` /
`:native_compiled` tags and stay gated behind the suite's own `:vit_full`
/ `:whisper_full` moduletag. `--only vit_full` now runs all three lanes
on the full checkpoint, while `--only native` stays tiny-random only and
never pulls a full-size download.

Both full forwards lower fully under `native_fallback: :raise` and hold
at the pinned 1e-4 reference tolerance in the fusion lane too —
`mx::compile` only reassociates f32, so its drift stays well below the
approximate `mx::fast::*` kernels (which need 1e-3 ViT / 1e-2 Whisper in
the `:fast_kernels_full` variants). The serving and fast-kernels tests
remain eval-only.
`Nx.take_along_axis` (the `Nx.Block.TakeAlongAxis` block) now lowers
under `compiler: Emily.Compiler, native: true` instead of raising. The
MLX op and eager NIF already existed (`Emily.Native.take_along_axis` →
`mlx::core::take_along_axis`, used by `Emily.Backend.native_take_along_axis/4`);
only the IR path was missing.

- ir.ex: add opcode `take_along_axis` (81) and a `lower_block/5` clause
  mirroring the eager twin — cast indices to s32, emit, coerce to output
  type — alongside the existing `Nx.Block.Take` lowering.
- opcodes.hpp: add `TakeAlongAxis = 81`, bump `kOpcodeCount`, and a
  dispatch case calling `mx::take_along_axis(in[0], in[1], axis, s)`.
- compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases
  (last axis, axis 0, a 3-D transformer-shaped gather).

This was the last op forcing a fallback in
`Bumblebee.Text.question_answering`'s answer-span gather, so the
DistilBERT-QA `Nx.Serving` forward now compiles fully native and fused.
Wire it into the conformance suite: extend `mode_test` with a `:tag`
option (stamp every lane) so the serving test runs its eval / native /
native_compiled lanes gated behind `:distilbert_full`.
The opcode wire values are hand-maintained in two places —
`Emily.IR`'s `@opcodes` map and the `Opcode` enum + `kOpcodeCount` in
`c_src/emily/opcodes.hpp` — with nothing enforcing they agree. A
mismatch compiles fine and only misbehaves at runtime.

Add `Emily.OpcodeParityTest`, which parses the header and asserts both
sides are a unique, gap-free `0..N-1` with N == `kOpcodeCount`. The
check is value-based, not name-based, so it doesn't depend on the
snake_case names matching the PascalCase enum (`negate`/`Negative`,
`fast_rms_norm`/`FastRMSNorm`); a name/value permutation that keeps both
contiguous is still caught by the equivalence suite. Expose
`Emily.IR.opcodes/0` for the test to read the map.
The example documents are Livebooks (`.livemd`), so name the directory
to match. Rename `notebooks/` → `livebooks/` (all 8 files) and update
every current-state reference:

- mix.exs: the `extras:` paths and the `groups_for_extras` regex, plus
  the nav-group label `Notebooks:` → `Livebooks:`.
- scripts/test-livebooks.sh: `NB_DIR` and the path comment.
- README.md / ROADMAP.md: the `livebooks/` path and the HexDocs
  "Livebooks" section references.

CHANGELOG.md and PLAN.md keep their `notebooks/` mentions — those are
point-in-time historical records, not current-structure docs.
@ausimian
ausimian changed the base branch from feat/expr-compiler-take-along-axis to feat/expr-compiler June 6, 2026 05:37
@ausimian ausimian closed this Jun 6, 2026
@ausimian ausimian reopened this Jun 6, 2026
@ausimian
ausimian merged commit 82ce9c6 into feat/expr-compiler Jun 6, 2026
3 checks passed
@ausimian
ausimian deleted the feat/expr-compiler-livebooks-rename branch June 6, 2026 05:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant